#define pexit(s) { fprintf( stderr, "%s\n", s ); exit(1); }
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include "jerasure.h"
#include "reed_sol.h"
using namespace std;

int main(int argc, char** argv)
{
	if( argc != 7 && argc != 6 )
		pexit( "usage: ./raid n k w packetsize raid_dir" );

	char *dir;
	char tempstring[100];
	long int offset, bytesread;
	int begin, end;
	int n, k, w, packetsize;
	int i;
	int first, print;
	JER_Slices *slices;
	unsigned char *buffer;
	vector<JER_Region*> sectorfailures;
	JER_Region* tempr;

	FILE **fps;

	srand48( 0 );
	n = atoi( argv[1] );
	k = atoi( argv[2] );
	w = atoi( argv[3] );
	packetsize = atoi( argv[4] );
	dir = argv[5];
	print = 0;

	fps = (FILE**) malloc( n * sizeof( FILE* ) );

	for( i = 0; i < n; i++ )
	{
		sprintf( tempstring, "%s/file%d", dir, i );
		if( i < k )
		fps[i] = fopen( tempstring, "r+" );
		if( fps[i] == NULL || i >= k )
			fps[i] = fopen( tempstring, "w+" );
	}
	offset = 0;
	first = 1;
	if( print )
		first = 0;
	fprintf( stdout, "Drive files opened.\n" );

	slices = new JER_Slices( n, k, packetsize, w, 
			RS_Classic_Generator( k, n - k, w ) );
	buffer = (unsigned char*) malloc( packetsize * w * n * sizeof(char) );
	for( i = 0; i < n; i++ )
	{
		slices->Ptrs[i] = buffer + i * packetsize * w * sizeof(char);
	}

	fprintf( stdout, "Performing encoding.\n" );
	while( !feof(fps[0]) )
	{
		if( first == 0 )
			fprintf( stdout, "Before encoding:\n" );
		for( i = 0; i < k; i++ )
		{
			fseek( fps[i], offset, SEEK_SET );
			if( i == 0 )
				bytesread = fread( slices->Ptrs[i], 1, packetsize * w, fps[i] );
			else 
			{
				if( bytesread != fread( slices->Ptrs[i], 1, packetsize * w, fps[i] ) )
				{
					fprintf( stdout, "Bytes read from file0 and file%d are not the same\n", i );
					exit(1);
				}
			}
			if( first == 0 )
			{
				fprintf( stdout, "Drive %d: ", i );
				fwrite( slices->Ptrs[i], 1, bytesread, stdout );
				fprintf( stdout, "\n" );
			}
		}
		if( !feof(fps[0]) || bytesread != 0 )
		{
			if( slices->Encode() )
				pexit( "Encoding failed." );
			if( first == 0 )
				fprintf( stdout, "After encoding:\n" );
			for( i = 0; i < n; i++ )
			{
				if( i < k )
					fseek( fps[i], offset, SEEK_SET );
				if( bytesread != fwrite( slices->Ptrs[i], 1, bytesread, fps[i] ) )
					pexit( "Does not read same amount of data from all files" );
				if( first == 0 )
				{
					fprintf( stdout, "Drive %d: ", i );
					fwrite( slices->Ptrs[i], 1, bytesread, stdout );
					fprintf( stdout, "\n" );
				}
			}
			offset += bytesread;
			first++;
		}
	}
	fprintf( stdout, "Encoding successful.\n" );

	fprintf( stdout, "Generating %d partial failures.\n", n - k );
	for( i = 0; i < n - k; i++ )
	{
		sectorfailures.push_back( new JER_Region );
		sectorfailures[i]->drive = lrand48() % n;
		sectorfailures[i]->start = lrand48() % offset;
		sectorfailures[i]->size = lrand48() % 
			( offset - sectorfailures[i]->start );
		fprintf( stdout, "Drive %d failed at %d for %d bytes\n",
				sectorfailures[i]->drive, sectorfailures[i]->start,
				sectorfailures[i]->size );
	}

	for( i = 0; i < n - k - 2; i++ )
	{
		first = lrand48() % k;
		fprintf( stdout, "Failing drive %d.\n", first );
		slices->States[first] = 1;
	}

	offset = 0;
	first = 1;
	if( print )
		first = 0;
	for( i = 0; i < n; i++ )
		fseek( fps[i], 0, SEEK_SET );

	fprintf( stdout, "Performing partial failure recovery and decoding.\n" );
	while( !feof(fps[0]) )
	{
		slices->Pfs.clear();
		if( first == 0 )
			fprintf( stdout, "Before decoding:\n" );
		for( i = 0; i < n; i++ )
		{
			fseek( fps[i], offset, SEEK_SET );
			if( i == 0 )
				bytesread = fread( slices->Ptrs[i], 1, packetsize * w, fps[i] );
			else
				fread( slices->Ptrs[i], 1, packetsize * w, fps[i] );
			bzero( slices->Ptrs[i] + bytesread, packetsize * w - bytesread );
			if( first == 0 )
			{
				fprintf( stdout, "Drive %d: ", i );
				fwrite( slices->Ptrs[i], 1, bytesread, stdout );
				fprintf( stdout, "\n" );
			}
		}
		sectorfailures.resize(0);
		for( i = 0; i < sectorfailures.size(); i++ )
		{
			tempr = sectorfailures[i];
			/* If any part of the sector failure is in this block, */
			if( tempr->start < bytesread + offset &&
					tempr->start + tempr->size > offset )
			{
				/* Add to Pfs a region beginning at tempr->start or beginning of 
				   packet, whichever comes later, and ends at the end of tempr or 
				   end of packet, whichever comes sooner. */
				begin = ( tempr->start < offset ) ? offset : tempr->start;
				end = ( tempr->start + tempr->size > offset + bytesread ) ? offset + bytesread - begin : tempr->start + tempr->size - begin;
				slices->Add_Partial_Failure( tempr->drive, begin - offset, end );

				/* Delete the region added to Pfs from the sectorfailures vector. 
				   Only delete the portion added to Pfs though. */

				/* If begin == start, no portion of the region exists before the 
				   beginning of the packet. The start of the region equals the
				   end of the packet and size is decremented by end, size of the
				   region added to Pfs. If this gives region a 0 or negative size
				   delete region from sectorfailures. */
				if( begin == tempr->start )
				{
					tempr->start = offset + bytesread;
					tempr->size -= end;
					if( tempr->size <= 0 )
					{
						sectorfailures.erase( sectorfailures.begin() + i );
						i--;
					}
				}
				/* Otherwise, if the end extends past the end of the packet,
				   create a new region. The original is the region that is 
				   before the current packet, which keeps the same start value
				   and size is start - begin. The new region is the portion
				   after the packet which has start at the end of the packet
				   and size of the old endpoint - the end of packet. This 
				   branch and the next should hopefully not be used. */
				else if( tempr->start + tempr->size > offset + bytesread )
				{
					sectorfailures.push_back( new JER_Region );
					sectorfailures[sectorfailures.size()-1]->drive = tempr->drive;
					sectorfailures[sectorfailures.size()-1]->start = 
						offset + bytesread;
					sectorfailures[sectorfailures.size()-1]->size = 
						tempr->start + tempr->size - offset - bytesread;
					tempr->size = begin - tempr->start;
				}
				/* Otherwise there is only a portion before the current packet.
				   Its start remains the same, but it is resized to begin - start. */
				else
				{
					tempr->size = begin - tempr->start;
				}
			}
		}

		if( !feof(fps[0]) || bytesread != 0 )
		{
			if( first == 0 )
			{
				for( i = 0; i < slices->Pfs.size(); i++ )
				{
					fprintf( stdout, "Partial failure on drive %d starting at %d with size %d being fixed.\n", slices->Pfs[i]->drive, slices->Pfs[i]->start, slices->Pfs[i]->size );
				}
			}
			if( slices->Decode() )
			{
				pexit( "Could not successfully decode." );
			}

			if( first == 0 )
				fprintf( stdout, "After decoding:\n" );
			for( i = 0; i < n; i++ )
			{
				fseek( fps[i], offset, SEEK_SET );
				fwrite( slices->Ptrs[i], 1, bytesread, fps[i] );
				if( first == 0 )
				{
					fprintf( stdout, "Drive %d: ", i );
					fwrite( slices->Ptrs[i], 1, bytesread, stdout );
					fprintf( stdout, "\n" );
				}
			}
			offset += bytesread;
			first++;
		}
	}

	fprintf( stdout, "Files successfully decoded.\n" );

	return 0;
}
